import rclpy
from rclpy.node import Node
from origincar_msg.msg import Sign
from sensor_msgs.msg import Image
import cv2
from cv_bridge import CvBridge
from pyzbar.pyzbar import decode
import numpy as np
import psutil
import time

class QrcodeNode(Node):
    def __init__(self, name):
        super().__init__(name)
        self.bridge = CvBridge()
        self.get_logger().info("大家好，我是发布二维码识别的话%s" % name)
        
        # 创建发布者，发布到 /sign_switch 话题
        self.publisher_ = self.create_publisher(Sign, 'sign_switch', 10)
        
        # 创建订阅者，订阅 /image 话题
        self.subscription = self.create_subscription(
            Image,
            '/image',
            self.image_callback,
            10
        )
        
        # 初始化微信二维码检测器
        self.wechat_qr_detector = QRCodeDetector()

    def image_callback(self, msg):
        try:
            # 将 ROS 压缩图像消息转换为字节数据
            np_arr = np.frombuffer(msg.data, np.uint8)
            
            # 使用 OpenCV 解码 JPEG 图像
            image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            image_roi = image[150:320,0:640]

            # 选择解码方法，注释掉不需要的部分以使用另一种方法
            self.decode_qr_code_pyzbar(image_roi)  # 使用 pyzbar 解码
            #self.decode_qr_code_wechat(image_roi)  # 使用 WeChatQRCode 解码

        except Exception as e:
            self.get_logger().error(f"Error processing image: {str(e)}")

    def decode_qr_code_pyzbar(self, image):
        # 使用 pyzbar 进行二维码解码
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        decoded_objects = decode(gray)
        
        # 遍历结果并判断是否是顺时针或逆时针二维码
        for obj in decoded_objects:
            qr_code_data = obj.data.decode('utf-8')
            if qr_code_data in ["ClockWise", "AntiClockWise"]:
                self.publish_sign_message(qr_code_data)
                return  # 成功识别并发布后立即返回，避免重复处理

    def decode_qr_code_wechat(self, image):
        # 使用微信二维码检测器进行解码
        qr_code_data = self.wechat_qr_detector.detect(image)
        
        # 判断是否是顺时针或逆时针二维码
        if qr_code_data in ["ClockWise", "AntiClockWise"]:
            self.publish_sign_message(qr_code_data)

    def publish_sign_message(self, qr_code_data):
        # 创建并发布 Sign 消息
        sign_msg = Sign()
        sign_msg.sign_data = 3 if qr_code_data == "ClockWise" else 4
        self.publisher_.publish(sign_msg)
        self.get_logger().info(f"已发布二维码信息: {qr_code_data}")

class QRCodeDetector:
    def __init__(self):
        self.depro = '/root/lyz_qrcode/detect.prototxt'
        self.decaf = '/root/lyz_qrcode/detect.caffemodel'
        self.srpro = '/root/lyz_qrcode/sr.prototxt'
        self.srcaf = '/root/lyz_qrcode/sr.caffemodel'

    def detect(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        detector = cv2.wechat_qrcode_WeChatQRCode(self.depro, self.decaf, self.srpro, self.srcaf)
        barcodes, _ = detector.detectAndDecode(gray)
        return barcodes[0] if barcodes else None

def main(args=None):
    rclpy.init(args=args)  # 初始化rclpy
    node = QrcodeNode("qrcode")  # 创建节点
    rclpy.spin(node)  # 保持节点运行
    rclpy.shutdown()  # 关闭rclpy

if __name__ == '__main__':
    main()
